from sonetto.agents.chat import ChatAgent
from sonetto.core.scheduler import Scheduler
import json

base_url = "https://open.bigmodel.cn/api/paas/v4"
api_key = "ff37fbba1e80d4cb5f38ac819893f4a1.c6T0Zm3Sd9YODRLV"
model = "glm-4-flash"

base_url_ds = 'https://api.siliconflow.cn/v1'
# base_url_ds = "https://api.deepseek.com"
# api_key_ds = "sk-3bb4d61950df4f4d833f800848e1f461"
api_key_ds = 'sk-djfheakeqfjnmvtpdlultealntivaqpgjmjxypefugifzgog'
model_ds = "deepseek-ai/DeepSeek-V3"

base_url_qwq = 'https://api-g6d5heobr54ac2p0.aistudio-app.com/v1'
api_key_qwq = '2584eed28ce01d54e919ec892f3f7ae4201e8069'
model_qwq = 'qwq'

scheduler = Scheduler()
writer = ChatAgent(
    "诗人",
    api_key,
    model,
    duty="写出漂亮的诗歌来通过评委的审核。",
    base_url=base_url,
    system_prompt="""你是一个诗人，你的任务是写出漂亮的诗歌来通过评委的审核。你写的诗要交给评委审核而非用户""",
    use_protocol=True,
)
scheduler.register_agent(writer)
evaluator = ChatAgent("评委",
                      api_key=api_key,
                      model=model,
                      system_prompt="""你是一个诗歌评委，你的任务是审核诗人写出的诗歌，并指出其中的不足。
写诗是一个学习的过程，所以评委不应该直接帮诗人改诗，而是要让他自己改进。
审核标准应当尽量高些。如果评委认为诗人的诗歌已经足够优秀，评委需要将 Sonetto 通信协议中的type置为“审核通过”。          
""",
                      duty="审核诗人写出的诗歌，并指出其中的不足。",
                      base_url=base_url,
                      use_protocol=True)
scheduler.register_agent(evaluator)
# scheduler.set_route(writer.agent_id, evaluator.agent_id)
# scheduler.set_route(evaluator.agent_id, writer.agent_id)

scheduler.run_stream("写一首赞美生命的诗歌", "审核通过", ["评委"], max_step=10)

# print(writer.sonetto_prompt)
# print(evaluator.sonetto_prompt)
# print([json.loads(ele['content']) for ele in writer.get_context()])
# print([json.loads(ele['content']) for ele in evaluator.get_context()])
